fun main() {
    val nums1 = intArrayOf(2, 0, 1, 2, 1, 0)
    val nums2 = intArrayOf(2, 0, 1)

    val s = Solution()
    s.sortColors(nums1)
    s.sortColors(nums2)

    println(nums1.contentToString())
    println(nums2.contentToString())
}

fun exchange(arr: IntArray, i: Int, j: Int) {
    val temp = arr[i]
    arr[i] = arr[j]
    arr[j] = temp
}

class Solution {
    fun sortColors(nums: IntArray): Unit {
        var i = 0
        var lessThan = 0
        var greatThan = nums.size

        while (i < greatThan) {
            if (nums[i] == 0) {
                exchange(nums, i, lessThan)
                lessThan += 1
                i += 1
            }else if (nums[i] == 1) {
                i += 1
            }else if (nums[i] == 2) {
                greatThan -= 1
                exchange(nums, i, greatThan)
            }
        }
    }
}